import json
import os
from functools import partial
from multiprocessing.pool import Pool
from pathlib import Path

import cv2
import pandas as pd
from tqdm import tqdm

from oulu_npu import retinex


def rgb_to_msr(rgb_path, msr_root, config, mode='train'):
    if mode == 'train':
        video_path = os.path.join(msr_root, 'Train_files_MSRCR', Path(rgb_path).parent.name)
    elif mode == 'dev':
        video_path = os.path.join(msr_root, 'Dev_files_MSRCR', Path(rgb_path).parent.name)
    else:
        video_path = os.path.join(msr_root, 'Test_files_MSRCR', Path(rgb_path).parent.name)

    if not os.path.exists(video_path):
        try:
            os.mkdir(video_path)
        except OSError:
            pass
    msr_path = os.path.join(video_path, Path(rgb_path).name)

    if os.path.exists(os.path.join(video_path, Path(rgb_path).name)):
        return

    # 遍历该目录下的所有图片文件
    # print("filename :", rgb_path)
    img = cv2.imread(rgb_path)

    # change to gray
    # （下面第一行是将RGB转成单通道灰度图，第二步是将单通道灰度图转成3通道灰度图）
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # print("img.shape: ", img.shape)
    img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    # print("img: ",img)
    img_MSRCR = retinex.MSRCR(
        img,
        config['sigma_list'],
        config['G'],
        config['b'],
        config['alpha'],
        config['beta'],
        config['low_clip'],
        config['high_clip']
    )
    # print("image_MSR.shape :", img_MSRCR.shape)
    # print("image_MSR :", img_MSRCR)



    # cv2.imwrite('D:\\Datasets\\test_2'+ "\\"+filename,img_MSRCR)
    cv2.imwrite(msr_path, img_MSRCR)


def get_images_paths(rgb_root, mode, data_csv):
    paths = []
    df = pd.read_csv(data_csv)
    for row in df.iterrows():
        if mode == 'train':
            path = os.path.join(rgb_root, 'Train_files_crops', row[1]['video'], row[1]['file'])
        elif mode == 'dev':
            path = os.path.join(rgb_root, 'Dev_files_crops', row[1]['video'], row[1]['file'])
        else:
            path = os.path.join(rgb_root, 'Test_files_crops', row[1]['video'], row[1]['file'])
        paths.append(path)
    return paths


def main():
    with open('config_msr.json', 'r') as f:
        config = json.load(f)
    # mode = 'train'
    # mode = 'dev'
    mode = 'test'
    rgb_root = '/home/shaohua/data2/Datasets/Face_Anti_Spoofing/Oulu_NPU'
    msr_root = '/home/shaohua/data2/Datasets/Face_Anti_Spoofing/Oulu_NPU'
    paths = get_images_paths(rgb_root, mode, '../data/data_{}.csv'.format(mode))
    with Pool(processes=2) as p:
        with tqdm(total=len(paths)) as pbar:
            for v in p.imap_unordered(partial(rgb_to_msr, msr_root=msr_root, config=config, mode=mode), paths):
                pbar.update()


if __name__ == "__main__":
    main()
